using System;
using System.Reflection;
using NBody;
using NBody.InitialConditions;
using NBody.Cosmology;
using OptionParser;
using Gsl;
using Fftw;
using DataFile;

[assembly: AssemblyTitle ("NBodyInit")]
[assembly: AssemblyVersion ("1.0.*")]
[assembly: AssemblyDescription ("Generate initial conditions for a single halo.")]
[assembly: AssemblyCopyright ("2005 Joseph D. MacMillan")]

public class InitOptions : OutputGetOptions
{
    [Option ("-n", "Number of particles per side of initial cube")]
    public int N = 32;
    
    [Option ("-l", "Radius of initial halo")]
    public double R = 186.0;
    
    [Option ("-z", "Initial redshift")]
    public double Z = 15.0;
    
    [Option ("-s", "Smoothing scale for the power spectrum")]
    public double SmoothingScale = 100.0;
    
    [Option ("-p", "Peak height for the density peak")]
    public double PeakHeight = 3.0;
    
    [Option ("-interp", "Number of points for interpolation")]
    public int NInterp = 100;
    
    [Option ("-density", "Only output density function")]
    public bool Density = false;
    
    [Option ("-perturb", "Apply small scale perturbations")]
    public bool Perturb = false;
    
    [Option ("-novel", "Don't apply velocity perturbations")]
    public bool NoVelPerturb = false;
    
    [Option ("-min", "Minimum wavelength")]
    public int MinWave = Int32.MinValue;
    
    [Option ("-max", "Maximum wavelength")]
    public int MaxWave = Int32.MaxValue;
    
    [Option ("-seed", "Random seed")]
    public int Seed = 0;
}

public class NBodyInit
{
    static double A;
    static double h = 0.5;
    static double H0 = 2.4108e-3; // for kpc, 10^8 M_sun, etc, with h = 0.5
    static double sigma_8 = 0.63;
    static double R0 = 8.0 / h;
    static double Z;
    static double Rstar, gamma, theta, xi_0;
    static double Smooth;
    static double peak;
    
    public static void Main(string[] args)
    {
        InitOptions opts = new InitOptions();
        opts.ProcessArgs(args);
        
        // set up units in comoving Mpc
        Z = opts.Z;
        double Length = 2.0 * opts.R / 1000.0 * (1.0 + Z);
        Smooth = opts.SmoothingScale / 1000.0;
        peak = opts.PeakHeight;
        
        // normalize power spectrum
        Gsl.Integration integ = new Gsl.Integration(1000);
        integ.Function = Sigma0;        
        integ.ExecuteQAGIU(0.0);
        double sigma2 = 1.0 / 2.0 / Math.PI / Math.PI * integ.Result;
        A = sigma_8 * sigma_8 / sigma2;
        //A = sigma_8 / sigma2;
        Console.Error.WriteLine("Power spectrum amplitude = {0}", A);
        // scale power spectrum normalization to redshift
        A *= 1.0 / (1.0 + Z) / (1.0 + Z);
        Console.Error.WriteLine("  Scaled amplitude = {0}", A);
        
        // compute some neccessary stuff for delta(r)
        double s0 = SpectralMoment(0);
        double s1 = SpectralMoment(1);
        double s2 = SpectralMoment(2);
        Rstar = Math.Sqrt(3.0) * s1 / s2;
        gamma = s1*s1 / s2 / s0;
        theta = (3.0 * (1.0 - gamma * gamma) + (1.216 - 0.9 * Math.Pow(gamma, 4.0)) * Math.Exp(-0.5 * gamma * Math.Pow(0.5 * peak * gamma, 2.0))) / (Math.Sqrt(3.0 * (1.0 - gamma * gamma) + 0.45 + Math.Pow(peak * gamma / 2.0, 2.0)) + 0.5 * peak * gamma);
        xi_0 = Xi(1e-5);
        Console.Error.WriteLine("  Rstar = {0}", Rstar);
        Console.Error.WriteLine("  gamma = {0}", gamma);
        Console.Error.WriteLine("  theta = {0}", theta);
        Console.Error.WriteLine("  Xi(0) = {0}", xi_0);
        
        if (opts.Density)
        {
            for (int i = 1; i < opts.NInterp; i++)
            {
                double rf = Math.Pow(10.0, -3.0 + (3.0 + Math.Log10(0.5 * Length)) * i / (opts.NInterp - 1.0));
                double r = rf * 1000.0 / (1.0 + Z);
                double delta = Delta(rf);
                double rho_b = 3.0 * H0 * H0 / 8.0 / Math.PI * Math.Pow(1.0 + Z, 3.0);
                double rho = rho_b * (1.0 + delta);
                Console.WriteLine("{0} {1} {2}", Math.Log10(r), Math.Log10(rho), Math.Log10(delta));
                //Console.WriteLine("{0} {1} {2}", r, rho, delta);
            }
            Environment.Exit(0);
        }
        
        // step 1: Generate data table of ri, rf values
        Console.Error.Write("Generating interpolation table ");
        Table table = new Table(opts.NInterp, 2, "Interpolation table for intial conditons");
        table[0, 0] = 0.0;
        table[0, 1] = 0.0;
        integ = new Integration(1000);
        integ.Function = new Function(DeltaInt);
        for (int i = 1; i < opts.NInterp; i++)
        {
            double rf = i / (opts.NInterp - 1.0) * 0.5 * Length;
            integ.ExecuteQAG(1e-7, rf, Integration.Key.GSL_INTEG_GAUSS15);
            double ri = Math.Pow(rf*rf*rf + 3.0 * integ.Result, 1.0/3.0);
            
            table[i, 0] = ri;
            table[i, 1] = rf;      
            
            if (i % (opts.NInterp / 50) == 0)
                Console.Error.Write(".");
        }
        Console.Error.WriteLine();
        
        // step 2:  Generate sphere of appropriate size
        // hmm, better re-do normalization for non-smoothing power spectrum
        Smooth = 0.0;
        integ = new Gsl.Integration(1000);
        integ.Function = Sigma0;        
        integ.ExecuteQAGIU(0.0);
        sigma2 = 1.0 / 2.0 / Math.PI / Math.PI * integ.Result;
        A = sigma_8 * sigma_8 / sigma2;
        A *= 1.0 / (1.0 + Z) / (1.0 + Z);
        
        Console.Error.WriteLine("Generating sphere ... ");
        RandomGaussian random;
        if (opts.Seed == 0)
            random = new RandomGaussian();
        else
            random = new RandomGaussian(opts.Seed);
        double delta_k = 2.0 * Math.PI / Length;
            
        Fftw.Complex[] fx = new Fftw.Complex[opts.N * opts.N * (opts.N/2 + 1)];
        Fftw.Complex[] fy = new Fftw.Complex[opts.N * opts.N * (opts.N/2 + 1)];
        Fftw.Complex[] fz = new Fftw.Complex[opts.N * opts.N * (opts.N/2 + 1)];
        
        double a, b;
        double kx, ky, kz;
        int c = 0;
        for (int i = 0; i < opts.N; i++)
        {
            kx = delta_k * i;
            for (int j = 0; j < opts.N; j++)
            {
                ky = delta_k * j;
                for (int k = 0; k < opts.N/2 + 1; k++)
                {
                    kz = delta_k * k;
                       
                    double kk = Math.Sqrt(kx*kx + ky*ky + kz*kz);
                        
                    double r1 = random.NextDouble();
                    double r2 = random.NextDouble();
                        
                    if (kk == 0.0)
                    {
                        a = 0.0;
                        b = 0.0;
                    }
                    else if (i < opts.MinWave || j < opts.MinWave || k < opts.MinWave ||
                             i > opts.MaxWave || j > opts.MaxWave || k > opts.MaxWave)
                    {
                        a = 0.0;
                        b = 0.0;
                    }
                    else
                    {
                        double p = P(kk);
                        a = Math.Sqrt(p) / kk / kk * r1;
                        b = Math.Sqrt(p) / kk / kk * r2;
                    }
                        
                    fx[c].Re = 0.5 * kx * b;
                    fx[c].Im = 0.5 * kx * a;
                      
                    fy[c].Re = 0.5 * ky * b;
                    fy[c].Im = 0.5 * ky * a;
                      
                    fz[c].Re = 0.5 * kz * b;
                    fz[c].Im = 0.5 * kz * a;
                        
                    c++;
                }
            }
        }    
            
        PlanComplexToReal3D planx = new PlanComplexToReal3D(opts.N, opts.N, opts.N, fx, Plan.FFTW_ESTIMATE);
        planx.Execute();
        double[] sx = planx.RealOutput;
            
        PlanComplexToReal3D plany = new PlanComplexToReal3D(opts.N, opts.N, opts.N, fy, Plan.FFTW_ESTIMATE);
        plany.Execute();
        double[] sy = plany.RealOutput;
            
        PlanComplexToReal3D planz = new PlanComplexToReal3D(opts.N, opts.N, opts.N, fz, Plan.FFTW_ESTIMATE);
        planz.Execute();
        double[] sz = planz.RealOutput;
        
        // Generate a few simple stats about the displacements
        int ntotal = opts.N * opts.N * opts.N;
            
        double avg = 0.0;
        double max = 0.0;
        for (int i = 0; i < ntotal; i++)
        {
            double dd = Math.Sqrt(sx[i]*sx[i] + sy[i]*sy[i] + sz[i]*sz[i]);
            avg += dd;
            if (dd > max) max = dd;
        }
        avg /= (double)ntotal;
        double grid = Length / (double)opts.N;
            
        Console.Error.WriteLine("Grid spacing: {0:f6}", grid);
        Console.Error.WriteLine("Avg displacements: {0:f6}", avg / (1.0 + Z));
        Console.Error.WriteLine("Max displacements: {0:f6}", max / (1.0 + Z));
        Console.Error.WriteLine("Redhsift should be around z = {0}", 1.0 / Math.Sqrt(grid / (1.0 + Z) / max) - 1.0);
        
        // aw, nuts.  Cutting out the few largest waves will reduce the amplitude of the 
        // perturbations quite a bit.  Better artificially scale them.
        double scale = grid / max * (1.0 + Z);
        
        // now apply displacements to particles on a cubic grid
        double time = 2.0 / 3.0 / H0 * Math.Pow(1.0 + Z, -3.0 / 2.0);
        NBodySystem s = new NBodySystem(ntotal, time);
        double delta_x = Length / (double)opts.N;
        double Mtot = 3.0 * H0 * H0 / 8.0 / Math.PI * Length * Length * Length * 1e9; // convert length units to kpc        
        double m = Mtot / (double)ntotal;
        c = 0;
        for (int i = 0; i < opts.N; i++)
        {
            double x = delta_x * i;
            for (int j = 0; j < opts.N; j++)
            {
                double y = delta_x * j;
                for (int k = 0; k < opts.N; k++)
                {
                    double z = delta_x * k;
                    
                    double xs, ys, zs;
                    double vx, vy, vz;
                    if (opts.Perturb)
                    {
                        xs = x + 1.0 / (1.0 + Z) * scale * sx[c];
                        while (xs > Length) xs -= Length; while (xs < 0.0) xs += Length;
                        ys = y + 1.0 / (1.0 + Z) * scale * sy[c];
                        while (ys > Length) ys -= Length; while (ys < 0.0) ys += Length;
                        zs = z + 1.0 / (1.0 + Z) * scale * sz[c];
                        while (zs > Length) zs -= Length; while (zs < 0.0) zs += Length;
                    
                        if (opts.NoVelPerturb)
                        {
                            vx = vy = vz = 0.0;
                        }
                        else
                        {
                            vx = H0 * Math.Sqrt(1.0 + Z) * scale * sx[c];
                            vy = H0 * Math.Sqrt(1.0 + Z) * scale * sy[c];
                            vz = H0 * Math.Sqrt(1.0 + Z) * scale * sz[c];
                        }
                    }
                    else
                    {
                        xs = x;
                        ys = y;
                        zs = z;
                        vx = vy = vz = 0.0;
                    }
                      
                    s.AddParticle(new Particle(m, xs, ys, zs, vx, vy, vz));
                    
                    c++;
                }
            }
        }
        
        s.AdjustForCentreOfMass();
        s.ExtractSphere(0.5 * Length);
        
        // apply hubble flow
        for (int i = 0; i < s.NumParts; i++)
        {
            double vx = s[i].Vx + 2.0 / 3.0 / time * s[i].X;
            double vy = s[i].Vy + 2.0 / 3.0 / time * s[i].Y;
            double vz = s[i].Vz + 2.0 / 3.0 / time * s[i].Z;
            s[i].SetVelocity(vx, vy, vz);
        }
        
        // step 3: perturb particles to form density profile
        Console.Error.WriteLine("Perturbing particles ... ");
        for (int i = 0; i < s.NumParts; i++)
        {
            double ri = s[i].Radius;
            double rf = table.Interpolate(0, 1, ri);
            
            if (ri != 0.0)
            {
                double x = s[i].X * rf / ri;
                double y = s[i].Y * rf / ri;
                double z = s[i].Z * rf / ri;
                
                double vx = s[i].Vx + 2.0 / 3.0 / time * (x - s[i].X);
                double vy = s[i].Vy + 2.0 / 3.0 / time * (y - s[i].Y);
                double vz = s[i].Vz + 2.0 / 3.0 / time * (z - s[i].Z);
                
                // now scale positions and velocites to correct (non-comoving, kpc) redshift
                double sc = 1000.0 / (1.0 + Z);
                s[i].SetPosition(sc * x, sc * y, sc * z);
                s[i].SetVelocity(sc * vx, sc * vy, sc * vz);
            }
        }
        
        s.Write(opts.OutFile);
    }
    
    static double Sigma0(double k, object o)
    {
        double q = k;
        double x = k * R0;
        double w = 3.0 * (Math.Sin(x) - x * Math.Cos(x)) / (x * x * x);
        double T = BBKS(q);
        double smooth = Math.Exp(-0.5 * k * k * Smooth * Smooth);
        return q*q*q * T * T * w * w * smooth;
    }
    
    static double P(double k)
    {
        double q = k;
        double T = BBKS(q);
        double smooth = Math.Exp(-0.5 * k * k * Smooth * Smooth);
        return A * q * T * T * smooth;
    }
            
    static double BBKS(double k)
    {
        double q = k / h / h;
        double a0 = 2.34 * q;
        double a1 = 3.89 * q;
        double a2 = 16.1 * q;
        double a3 = 5.46 * q;
        double a4 = 6.71 * q;
        double T = Math.Log(1.0 + a0) / a0 / Math.Sqrt(Math.Sqrt(1.0 + a1 + a2*a2 + a3*a3*a3 + a4*a4*a4*a4));
            
        return T;
    }
    
    static double Xi(double r)
    {
        Gsl.Integration i = new Gsl.Integration(1000);
        i.Function = new Gsl.Function(XiInt);
        i.Parameters = r;
        i.ExecuteQAGIU(0.0);
            
        double xi = 1.0 / 2.0 / Math.PI / Math.PI / r * i.Result;
        return xi;
    }
    
    static double XiInt(double k, object radius)
    {
        double r = (double)radius;
        double p = P(k);
        double i = p * k * Math.Sin(k * r);
        //Console.WriteLine("{0} {1}", k, i);
        return i;
    }
    
    static double SpectralMoment(int l)
    {
        Gsl.Integration i = new Gsl.Integration(1000);
        i.Function = new Gsl.Function(SpectralMomentInt);
        i.Parameters = l;
        i.ExecuteQAGIU(0.0);
            
        double ans = 1.0 / 2.0 / Math.PI / Math.PI * i.Result;
        return Math.Sqrt(ans);
    }
        
    static double SpectralMomentInt(double k, object o)
    {
        double l = (int)o;
        return P(k) * Math.Pow(k, 2.0 * (l+1.0));
    }
    
    static double CorrFuncGrad2(double r)
    {
        Gsl.Integration i = new Gsl.Integration(1000);
        i.Function = new Gsl.Function(CorrFuncGrad2Int);
        i.Parameters = r;
        i.ExecuteQAGIU(0.0);
            
        double ans = -1.0 / (6.0 * Math.PI * Math.PI * r) * i.Result;
        return ans;
    }
        
    static double CorrFuncGrad2Int(double k, object o)
    {
        double r = (double)o;
        double p = P(k);
        double y = k*k* Math.Sin(k * r);
        double i = p * k * y;
        return i;
    }
    
    static double Delta(double x)
    {
        double r = x;
        double xi = Xi(r);
        double xi_diff2 = CorrFuncGrad2(r);
        double d = peak * xi / Math.Sqrt(xi_0) - theta / gamma / (1.0 - gamma * gamma) * (gamma*gamma * xi + Rstar * Rstar * xi_diff2) / Math.Sqrt(xi_0);
        return d;
    }
    
    static double DeltaInt(double x, object o)
    {
        return Delta(x) * x * x;
    }
}

public class RandomGaussian : Random
    {
                   
        private static bool have_next;
        private static double next;
        
        public RandomGaussian() : base()
        {
            next = 0.0;
            have_next = false;
        }
        
        public RandomGaussian(int seed) : base(seed)
        {
            next = 0.0;
            have_next = false;
        }
        
        protected override double Sample()
        {
            if (have_next)
            {
                have_next = false;
                return next;
            }
            else
            {
                double x = 0.0;
                double y = 0.0;
                while (x == 0.0 || y == 0.0)
                {
                    x = base.Sample();
                    y = base.Sample();
                }
                next = Math.Sqrt(-2.0 * Math.Log(x)) * Math.Sin(2.0 * Math.PI * y);
                
                have_next = true;
                return Math.Sqrt(-2.0 * Math.Log(x)) * Math.Cos(2.0 * Math.PI * y);
            }
        }
    }
